#!/usr/bin/python3
# -*- coding: utf-8 -*-

# *************************************************************************
#    > File Name: client.py
#    > Author: xlzh
#    > Mail: xiaolongzhang2015@163.com
#    > Created Time: 2025年01月17日 星期五 14时19分36秒
# *************************************************************************

import sys
import requests
from requests.exceptions import ConnectionError
import subprocess
from settings import *


def get_available_list():
    """
    get all available sample list from server
    """
    try:
        res = requests.get(PUBLIC_LIST)
    except ConnectionError as err:
        sys.stderr.write(f"Error: {err}\n")
        return

    if res.status_code != 200:
        sys.stderr.write("Error: failed to get public available list!\n")
        return

    data_list = res.json()['tasks']
    sys.stdout.write("AccessCode\tRetrievedTimes\tFileName\n")

    for item in data_list:
        sys.stdout.write(f"{item['access_code']}\t{item['species']}\t{item['retrieve_times']}\t{item['file_name']}\n")


def align_sequence_core(input_stream, out_file):
    """
    align the sequence and write to out SAM file
    """
    out_handle = open(out_file, "wb")
    bwa_command = [BWA_PATH, 'mem', '-t', str(BWA_THREAD), REF_PATH, '-']

    bwa_process = subprocess.Popen(
        bwa_command,
        stdin=subprocess.PIPE,
        stdout=out_handle,
        stderr=subprocess.PIPE
    )

    # put the fastq chunk into stdin
    for chunk in input_stream:
        bwa_process.stdin.write(chunk)

    # close the stdin after all data transfer to bwa_process
    bwa_process.stdin.close()
    bwa_process.wait()  # wait the bwa_process

    # close the SAM file handle
    out_handle.close()
    return bwa_process.returncode


def align_sequence(access_code: str, out_file: str):
    """
    align the sequence into bam file
    """
    try:
        res = requests.get(PUBLIC_ACCESS + f'?access_code={access_code}', stream=True)
    except ConnectionError as err:
        sys.stderr.write(f"Error: {err}\n")
        return

    if res.status_code != 200:
        sys.stderr.write(f"{res.json()['detail']}")
        return

    # open the stream to align and write to output SAM file
    align_sequence_core(res.iter_content(chunk_size=None), out_file)


def main():
    args = sys.argv

    if len(args) < 2:
        sys.stderr.write("usage: python3 client.py <BWA|LIST> [access_code] [out_sam_file]\n")
        sys.exit(-1)

    if args[1] == 'LIST':
        return get_available_list()

    elif args[1] == 'BWA':
        align_sequence(args[2], args[3])


if __name__ == '__main__':
    main()
